#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import division

from math import exp, log, sqrt
from scipy.stats import norm
from scipy import isscalar, asarray

def d_j(j, S, K, r, v, T):
    """
   d_j = \frac{log(\frac{S}{K})+(r+(-1)^{j-1} \frac{1}{2}v^2)T}{v sqrt(T)}
   """

    return (log(S / K) + (r + (-1) ** (j - 1) * 0.5 * v * v) * T) / (v
            * T ** 0.5)

def vanillaCallPrice(S, K, r, v, T):
    """
   Price of a European call option struck at K, with spot S,
   constant rate r, constant vol v (over the life of the option)
   and time to maturity T
   """
    if isscalar(S):
        return S * norm.cdf(d_j(1., S, K, r, v, T)) - \
               K * exp(-r * T) * norm.cdf(d_j(2, S, K, r, v, T))
    else:
        return asarray([ Si * norm.cdf(d_j(1., Si, K, r, v, T)) - \
               K * exp(-r * T) * norm.cdf(d_j(2, Si, K, r, v, T)) \
               for Si in S])

def vanillaPutPrice(S, K, r, v, T):
    """
     Price of a European put option struck at K, with spot S,
     constant rate r, constant vol v (over the life of the option)
     and time to maturity T
   """

    if isscalar(S):
        return -S * norm.cdf(-d_j( 1, S, K, r, v, T)) + K * exp(-r * T) \
                   * norm.cdf(-d_j(2, S, K, r, v, T))
    else:
        return asarray([ -Si * norm.cdf(-d_j( 1, Si, K, r, v, T)) + \
                K * exp(-r * T) * norm.cdf(-d_j(2, Si, K, r, v, T)) \
                for Si in S])

def vegaVanilla(S, K, r, v, T):
    """
        Put and Call have the same Vega, Put-Call paraity
    """
    if isscalar(S):
        return S * norm.cdf(d_j( 1, S, K, r, v, T)) * sqrt(T)
    else:
        return asarray([Si * norm.cdf(d_j( 1, Si, K, r, v, T)) \
                        for Si in S]) * sqrt(T)

# wikipidea
def cashOrNothingCall(S, K, r, v, T):
    S = asarray(S)
    tmp = exp(-r * T) * asarray([norm.cdf(d_j(2, Si, K, r, v, T)) \
            for Si in S], 'float')
    if 1 == S.size:
        return float(tmp)
    else:
        return tmp

def cashOrNothingPut(S, K, r, v, T):
    S = asarray(S)
    tmp = exp(-r * T) * asarray([norm.cdf(-d_j(2, Si, K, r, v, T)) \
            for Si in S], 'float')
    if 1 == S.size:
        return float(tmp)
    else:
        return tmp

def assetOrNothingCall(S, K, r, q, v, T):
        return S * exp( -q * T ) * norm.cdf(d_j(1, S, K, r, v, T))

def assethOrNothingPut(S, K, r, q, v, T):
        return S * exp( -q * T ) * norm.cdf(-d_j(1, S, K, r, v, T))

